package org.hong.monkey.network.server;

import com.google.common.base.Preconditions;
import io.netty.channel.Channel;
import org.hong.monkey.network.buffer.ManagedBuffer;
import org.hong.monkey.network.client.TransportClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Iterator;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

public class OneForOneStreamManager extends StreamManager {
    private final Logger logger = LoggerFactory.getLogger(OneForOneStreamManager.class);

    private final AtomicLong nextStreamId;
    private final ConcurrentHashMap<Long, StreamState> streams;

    private static class StreamState {
        final String appId;
        final Iterator<ManagedBuffer> buffers;

        Channel associatedChannel = null;

        int curChunk = 0;

        StreamState(String appId, Iterator<ManagedBuffer> buffers) {
            this.appId = appId;
            this.buffers = Preconditions.checkNotNull(buffers);
        }
    }

    public OneForOneStreamManager() {
        nextStreamId = new AtomicLong((long) new Random().nextInt(Integer.MAX_VALUE) * 1000);
        streams = new ConcurrentHashMap<Long, StreamState>();
    }

    @Override
    public void registerChannel(Channel channel, long streamId) {
        if (streams.containsKey(streamId)) {
            streams.get(streamId).associatedChannel = channel;
        }
    }

    @Override
    public ManagedBuffer getChunk(long streamId, int chunkIndex) {
        StreamState state = streams.get(streamId);
        if (chunkIndex != state.curChunk) {
            throw new IllegalStateException(String.format(
                    "Received out-of-order chunk index %s (expected %s)", chunkIndex, state.curChunk));
        } else if (!state.buffers.hasNext()) {
            throw new IllegalStateException(String.format(
                    "Requested chunk index beyond end %s", chunkIndex));
        }
        state.curChunk += 1;
        ManagedBuffer nextChunk = state.buffers.next();

        if (!state.buffers.hasNext()) {
            logger.trace("Removing stream id {}", streamId);
            streams.remove(streamId);
        }

        return nextChunk;
    }

    @Override
    public void connectionTerminated(Channel channel) {
        for (Map.Entry<Long, StreamState> entry: streams.entrySet()) {
            StreamState state = entry.getValue();
            if (state.associatedChannel == channel) {
                streams.remove(entry.getKey());

                while (state.buffers.hasNext()) {
                    state.buffers.next().release();
                }
            }
        }
    }

    @Override
    public void checkAuthorization(TransportClient client, long streamId) {
        if (client.getClientId() != null) {
            StreamState state = streams.get(streamId);
            Preconditions.checkArgument(state != null, "Unknown stream ID.");
            if (!client.getClientId().equals(state.appId)) {
                throw new SecurityException(String.format(
                        "Client %s not authorized to read stream %d (app %s).",
                        client.getClientId(),
                        streamId,
                        state.appId));
            }
        }
    }

    public long registerStream(String appId, Iterator<ManagedBuffer> buffers) {
        long myStreamId = nextStreamId.getAndIncrement();
        streams.put(myStreamId, new StreamState(appId, buffers));
        return myStreamId;
    }
}
